// Author: Zhang Yisong(zhang_yisong@sonosemi.com)
#pragma once

#include "pw_async2/coro.h"
#include "pw_async2/dispatcher.h"
#include "pw_async2/pend_func_task.h"
#include "pw_async2/poll.h"
#include "pw_async2/waker.h"
#include "pw_result/result.h"
#include <tuple>

namespace coro {

template <typename... Pendables>
class WhenAny {
private:
  static constexpr auto kTupleIndexSequence =
      std::make_index_sequence<sizeof...(Pendables)>();

public:
  explicit WhenAny(Pendables&&... pendables)
      : m_pendables(std::forward<Pendables>(pendables)...),
        m_outputs(pw::async2::Poll<pw::async2::PendOutputOf<Pendables>>(
            pw::async2::Pending())...) {}

  auto Pend(pw::async2::Context& cx) 
      -> pw::async2::Poll<std::tuple<std::optional<pw::async2::PendOutputOf<Pendables>>...>> {
    if (auto result = CheckCompleted(cx, kTupleIndexSequence)) {
      return pw::async2::Ready(std::move(*result));
    }
    return pw::async2::Pending();
  }

private:
  template <size_t... Is>
  std::optional<std::tuple<std::optional<pw::async2::PendOutputOf<Pendables>>...>>
  CheckCompleted(pw::async2::Context& cx, std::index_sequence<Is...>) {
    // 初始化全nullopt的tuple
    std::tuple<std::optional<pw::async2::PendOutputOf<Pendables>>...> result;
    bool any_ready = false;

    // 检查每个Pendable
    ([&] {
      auto& output = std::get<Is>(m_outputs);
      if (!output.IsReady()) {
        output = std::get<Is>(m_pendables).Pend(cx);
      }

      if (output.IsReady()) {
        any_ready = true;
        std::get<Is>(result).emplace(std::move(output).value());
      }
    }(), ...);

    return any_ready ? std::make_optional(std::move(result)) : std::nullopt;
  }

  std::tuple<Pendables...> m_pendables;
  std::tuple<pw::async2::Poll<pw::async2::PendOutputOf<Pendables>>...> m_outputs;
};


template <typename... Pendables>
WhenAny(Pendables &&...) -> WhenAny<Pendables...>;

} // namespace coro